# Copyright 2020 - 2022 MONAI Consortium
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#     http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from sklearn.model_selection import KFold  ## K折交叉验证

import os
import json
import math
import numpy as np
import torch
from monai import transforms, data
import SimpleITK as sitk
from tqdm import tqdm 
from torch.utils.data import Dataset, DataLoader
from .multiprocess_loader import get_multiprocess_dataloader
from .transformer import RandomRotate, \
    GammaTransform, RemoveLabelTransform, MirrorTransform, Compose, NumpyToTensor
from .transformer.crop import RandCropByPosNegLabel
from .transformer.noisy_transformers import GaussianNoiseTransform, GaussianBlurTransform
from .transformer.color_transformers import ContrastAugmentationTransform, BrightnessMultiplicativeTransform
from .transformer.resample_transformers import SimulateLowResolutionTransform
from monai import transforms
import h5py

train_transform = [
    RandomRotate(angle_spectrum=30, execution_probability=0.2),
    RemoveLabelTransform(-1, 0),
    # RandCropByPosNegLabel([96, 96, 96], pos=1.0, neg=1.0, num_samples=2),
    GaussianNoiseTransform(noise_variance=(0, 0.1), execution_probability=0.1),
    GaussianBlurTransform(blur_sigma=(0.5, 1), execution_probability=0.1, p_per_channel=0.5, different_sigma_per_channel=True),
    BrightnessMultiplicativeTransform(multiplier_range=(0.75, 1.25), execution_probability=0.15),
    ContrastAugmentationTransform(contrast_range=(0.75, 1.25), execution_probability=0.15),
    SimulateLowResolutionTransform(zoom_range=(0.5, 1.0), per_channel=True, p_per_channel=0.5,
                                   order_downsample=0, order_upsample=3, execution_probability=0.1),
    GammaTransform((0.7, 1.5), True, True, execution_probability=0.1),
    GammaTransform((0.7, 1.5), False, True, execution_probability=0.1),
    # MirrorTransform((0, 1, 2)),
    NumpyToTensor(),
]

train_transform = Compose(train_transform)

val_transform = [
    RemoveLabelTransform(-1, 0),
    NumpyToTensor()
]
val_transform = Compose(val_transform)

class Dataset3d(Dataset):
    """
    """
    def __init__(
        self,
        paths
    ) -> None:

        super(Dataset3d, self).__init__()
        self.cache_num = len(paths)
        if self.cache_num > 0:
            self._cache_image = [None] * self.cache_num

            self._cache_label = [None] * self.cache_num

            for i in range(self.cache_num):
                self._cache_image[i], self._cache_label[i] = \
                        self._load_cache_item(paths[i])

    def get_labels(self, label):
        labels = np.zeros(label.shape[1:])
        labels[label[0] == 1] = 1
        # labels[label[1] == 1] = 2
        return labels

    def _load_cache_item(self, d_path):
        h5_image = h5py.File(d_path, "r")
        image = h5_image["image"][()]
        label = h5_image["label"][()]
        h5_image.close()
        # print(image.shape)
        image = image[0:1] # 单模态
        # print(image.shape)

        labels = self.get_labels(label)

        return image, labels

    def __getitem__(self, index):
        image = self._cache_image[index]
        image_mean = image.mean()
        image_std = image.std()
        image = (image - image_mean) / image_std

        label = self._cache_label[index]

        return image.astype(np.float32), label.astype(np.float32)

    def __len__(self):
        return len(self._cache_image)


def get_huanhu_loader(train_path, val_path):

    train_ds = Dataset3d(train_path)

    train_loader = get_multiprocess_dataloader(train_ds, train_transform=train_transform, batch_size=1, num_process=8)
    
    val_ds = Dataset3d(val_path)
    
    val_loader = DataLoader(val_ds, batch_size=1, shuffle=False)

    loader = [train_loader, val_loader]

    return loader
